{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0_xya1nyDHfY",
   "metadata": {
    "id": "0_xya1nyDHfY"
   },
   "source": [
    "<table style=\"width:100%\">\n",
    "<tr>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<font size=\"2\">\n",
    "Supplementary code for the <a href=\"http://mng.bz/orYv\">Build a Large Language Model From Scratch</a> book by <a href=\"https://sebastianraschka.com\">Sebastian Raschka</a><br>\n",
    "<br>Code repository: <a href=\"https://github.com/rasbt/LLMs-from-scratch\">https://github.com/rasbt/LLMs-from-scratch</a>\n",
    "<br>汉化的库: <a href=\"https://github.com/GoatCsu/CN-LLMs-from-scratch.git\">https://github.com/GoatCsu/CN-LLMs-from-scratch.git</a>\n",
    "</font>\n",
    "</td>\n",
    "<td style=\"vertical-align:middle; text-align:left;\">\n",
    "<a href=\"http://mng.bz/orYv\"><img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/cover-small.webp\" width=\"100px\"></a>\n",
    "</td>\n",
    "</tr>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "l62zIRRSBy_R",
   "metadata": {
    "id": "l62zIRRSBy_R"
   },
   "source": [
    "# 把From-Scratch GPT 转换为Llama 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aFmxTQbwCUMl",
   "metadata": {
    "id": "aFmxTQbwCUMl"
   },
   "source": [
    "## 从零构建 Llama 3.2（独立笔记本）\n",
    "\n",
    "### 本笔记本内容\n",
    "在本笔记本中，我们将 **逐步将原始 GPT 架构转换为 Llama 2 模型**（值得注意的是，GPT 和 GPT-2 共享相同的架构）。\n",
    "\n",
    "### 为什么不是 Llama 1 或 Llama 3？\n",
    "- **Llama 1 vs. Llama 2**：\n",
    "  - Llama 1 的架构与 Llama 2 类似，但 **Llama 2 具有更大的上下文窗口**，这一点非常重要。\n",
    "  - Llama 1 的权重 **未公开**，且使用受限，因此更合理的选择是专注于 Llama 2。\n",
    "\n",
    "- **Llama 3**：\n",
    "  - 我将在另一个笔记本中 **演示如何从 Llama 2 迁移到 Llama 3**，其变动较小。\n",
    "\n",
    "### 说明\n",
    "- 本笔记本的解释内容**刻意简化**，以避免冗余，**重点关注代码实现**。\n",
    "- 更多详细信息请参考 **Llama 2 论文**：[Llama 2: Open Foundation and Fine-Tuned Chat Models (2023)](https://arxiv.org/abs/2307.09288)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ohhMKUWvGm9z",
   "metadata": {
    "id": "ohhMKUWvGm9z"
   },
   "source": [
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/gpt-to-llama/gpt2-to-llama2-llama3.webp?1\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "JBpQwU89ETA1",
   "metadata": {
    "id": "JBpQwU89ETA1"
   },
   "source": [
    "- 下面是所需要的库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "34a9a440-84c2-42cc-808b-38677cb6af8a",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "34a9a440-84c2-42cc-808b-38677cb6af8a",
    "outputId": "8118963b-3c72-43af-874b-439ffebdc94c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface_hub version: 0.24.7\n",
      "sentencepiece version: 0.2.0\n",
      "torch version: 2.4.1+cu121\n"
     ]
    }
   ],
   "source": [
    "from importlib.metadata import version\n",
    "\n",
    "pkgs = [\n",
    "    \"huggingface_hub\",  # to download pretrained weights\n",
    "    \"sentencepiece\",    # to implement the tokenizer\n",
    "    \"torch\",            # to implement the model\n",
    "]\n",
    "for p in pkgs:\n",
    "    print(f\"{p} version: {version(p)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "UJJneXpTEg4W",
   "metadata": {
    "id": "UJJneXpTEg4W"
   },
   "source": [
    "&nbsp;\n",
    "# 1. 逐步转换 GPT 模型实现"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "v1zpfX2GHBKa",
   "metadata": {
    "id": "v1zpfX2GHBKa"
   },
   "source": [
    "在本节中，我们将基于 [第 4 章](../../ch04/01_main-chapter-code/ch04.ipynb) 的 GPT 模型代码，并 **逐步修改** 以实现 Llama 2 的架构。\n",
    "\n",
    "随后，我们将加载 Meta AI **公开分享的 Llama 2 预训练权重**。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "979c7b6d-1370-4da1-8bfb-a2b27537bf2f",
   "metadata": {
    "id": "979c7b6d-1370-4da1-8bfb-a2b27537bf2f"
   },
   "source": [
    "&nbsp;\n",
    "## 1.1 替换 LayerNorm 为 RMSNorm 层"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8b27fc8-23a1-4e0e-a1ea-792e0428e5e6",
   "metadata": {
    "id": "f8b27fc8-23a1-4e0e-a1ea-792e0428e5e6"
   },
   "source": [
    "首先，我们将 **LayerNorm** 替换为 **均方根层归一化（RMSNorm）**。\n",
    "\n",
    "- **LayerNorm** 使用 **均值和方差** 归一化输入，而 **RMSNorm** 仅使用 **均方根（RMS）**，从而提高计算效率。\n",
    "- RMSNorm 计算公式如下，其中 $x$ 为输入，$\\gamma$ 为可训练参数（向量），$\\epsilon$ 为小常数，用于避免除零错误：\n",
    "\n",
    "$$y_i = \\frac{x_i}{\\text{RMS}(x)} \\gamma_i, \\quad \\text{其中} \\quad \\text{RMS}(x) = \\sqrt{\\epsilon + \\frac{1}{n} \\sum x_i^2}$$\n",
    "\n",
    "更多细节请参考论文：[Root Mean Square Layer Normalization (2019)](https://arxiv.org/abs/1910.07467)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d7094381-9499-4e9e-93f9-b79470da3771",
   "metadata": {
    "id": "d7094381-9499-4e9e-93f9-b79470da3771"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "#####################################\n",
    "# Chapter 4\n",
    "#####################################\n",
    "\n",
    "# class LayerNorm(nn.Module):\n",
    "#     def __init__(self, emb_dim):\n",
    "#         super().__init__()\n",
    "#         self.eps = 1e-5\n",
    "#         self.scale = nn.Parameter(torch.ones(emb_dim))\n",
    "#         self.shift = nn.Parameter(torch.zeros(emb_dim))\n",
    "\n",
    "#     def forward(self, x):\n",
    "#         mean = x.mean(dim=-1, keepdim=True)\n",
    "#         var = x.var(dim=-1, keepdim=True, unbiased=False)\n",
    "#         norm_x = (x - mean) / torch.sqrt(var + self.eps)\n",
    "#         return self.scale * norm_x + self.shift\n",
    "\n",
    "\n",
    "class RMSNorm(nn.Module):\n",
    "    def __init__(self, emb_dim, eps=1e-5):\n",
    "        super().__init__()\n",
    "        self.eps = eps\n",
    "        self.emb_dim = emb_dim\n",
    "        self.weight = nn.Parameter(torch.ones(emb_dim)).float()\n",
    "\n",
    "    def forward(self, x):\n",
    "        means = x.pow(2).mean(dim=-1, keepdim=True)\n",
    "        x_normed = x * torch.rsqrt(means + self.eps)\n",
    "        return (x_normed * self.weight).to(dtype=x.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mtWC8DOmIu0F",
   "metadata": {
    "id": "mtWC8DOmIu0F"
   },
   "source": [
    "以下代码单元用于检查该实现是否与 PyTorch 内置的 RMSNorm："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e41ade7a-bf06-48b1-8b7e-0e4037d5753f",
   "metadata": {
    "id": "e41ade7a-bf06-48b1-8b7e-0e4037d5753f"
   },
   "outputs": [],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "example_batch = torch.randn(2, 3, 4)\n",
    "\n",
    "rms_norm = RMSNorm(emb_dim=example_batch.shape[-1])\n",
    "rmsnorm_pytorch = torch.nn.RMSNorm(example_batch.shape[-1], eps=1e-5)\n",
    "\n",
    "assert torch.allclose(rms_norm(example_batch), rmsnorm_pytorch(example_batch))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5eb81f83-c38c-46a4-b763-aa630a32e357",
   "metadata": {
    "id": "5eb81f83-c38c-46a4-b763-aa630a32e357"
   },
   "source": [
    "&nbsp;\n",
    "## 1.2 用SiLU代替GELU作激活函数"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b8aa702-f118-4ff6-9135-90725ec8756c",
   "metadata": {
    "id": "0b8aa702-f118-4ff6-9135-90725ec8756c"
   },
   "source": [
    "Llama 采用 **SiLU（Sigmoid Linear Unit）** 作为激活函数，而非 GELU。SiLU 也被称为 **Swish 函数**，其计算公式如下：\n",
    "\n",
    "$$\n",
    "\\text{silu}(x) = x \\cdot \\sigma(x), \\quad \\text{其中} \\quad \\sigma(x) \\text{ 为 logistic sigmoid 函数。}\n",
    "$$\n",
    "\n",
    "更多信息请参考论文：[Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning (2017)](https://arxiv.org/abs/1702.03118)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a74f3757-c634-4a3a-a8f3-6334cde454fe",
   "metadata": {
    "id": "a74f3757-c634-4a3a-a8f3-6334cde454fe"
   },
   "outputs": [],
   "source": [
    "#####################################\n",
    "# Chapter 4\n",
    "#####################################\n",
    "\n",
    "# class GELU(nn.Module):\n",
    "#     def __init__(self):\n",
    "#         super().__init__()\n",
    "\n",
    "#     def forward(self, x):\n",
    "#         return 0.5 * x * (1 + torch.tanh(\n",
    "#             torch.sqrt(torch.tensor(2.0 / torch.pi)) *\n",
    "#             (x + 0.044715 * torch.pow(x, 3))\n",
    "#         ))\n",
    "\n",
    "\n",
    "class SiLU(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SiLU, self).__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x * torch.sigmoid(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "72ecbe2e-b6b7-4319-972b-1a7fefa3368c",
   "metadata": {
    "id": "72ecbe2e-b6b7-4319-972b-1a7fefa3368c"
   },
   "outputs": [],
   "source": [
    "silu = SiLU()\n",
    "\n",
    "assert torch.allclose(silu(example_batch), torch.nn.functional.silu(example_batch))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f9b5167-1da9-46c8-9964-8036b3b1deb9",
   "metadata": {
    "id": "4f9b5167-1da9-46c8-9964-8036b3b1deb9"
   },
   "source": [
    "&nbsp;\n",
    "## 1.3 更新 FeedForward 模块"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a381e7a-b807-472e-91c9-3e4e3fc5ad91",
   "metadata": {
    "id": "3a381e7a-b807-472e-91c9-3e4e3fc5ad91"
   },
   "source": [
    "实际上，Llama 使用的是 SiLU 的 **门控线性单元（Gated Linear Unit, GLU）** 变体，称为 **SwiGLU**。这使得 `FeedForward` 模块的结构略有不同。\n",
    "\n",
    "SwiGLU 在前馈层中引入了 **门控机制**，其计算公式如下：\n",
    "\n",
    "$$\\text{SwiGLU}(x) = \\text{SiLU}(\\text{Linear}_1(x)) * (\\text{Linear}_2(x))$$\n",
    "\n",
    "其中：\n",
    "- $\\text{Linear}_1$ 和 $\\text{Linear}_2$ 为两个线性变换层\n",
    "- $*$ 表示 **按元素（element-wise）乘法**\n",
    "- **第三个线性层** $\\text{Linear}_3$ 在该门控激活之后应用\n",
    "\n",
    "更多信息请参考论文：[GLU Variants Improve Transformer (2020)](https://arxiv.org/abs/2002.05202)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d25fbe3d-b7c9-4772-ad67-bc0527e4e20a",
   "metadata": {
    "id": "d25fbe3d-b7c9-4772-ad67-bc0527e4e20a"
   },
   "outputs": [],
   "source": [
    "#####################################\n",
    "# Chapter 4\n",
    "#####################################\n",
    "# class FeedForward(nn.Module):\n",
    "#     def __init__(self, cfg):\n",
    "#         super().__init__()\n",
    "#         self.layers = nn.Sequential(\n",
    "#             nn.Linear(cfg[\"emb_dim\"], 4 * cfg[\"emb_dim\"]),\n",
    "#             GELU(),\n",
    "#             nn.Linear(4 * cfg[\"emb_dim\"], cfg[\"emb_dim\"]),\n",
    "#         )\n",
    "\n",
    "#     def forward(self, x):\n",
    "#         return self.layers(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "477568cb-03cd-4510-b663-a42ce3ec64a2",
   "metadata": {
    "id": "477568cb-03cd-4510-b663-a42ce3ec64a2"
   },
   "outputs": [],
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.fc1 = nn.Linear(cfg[\"emb_dim\"], cfg[\"hidden_dim\"], dtype=cfg[\"dtype\"], bias=False)\n",
    "        self.fc2 = nn.Linear(cfg[\"emb_dim\"], cfg[\"hidden_dim\"], dtype=cfg[\"dtype\"], bias=False)\n",
    "        self.fc3 = nn.Linear(cfg[\"hidden_dim\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"], bias=False)\n",
    "        self.silu = SiLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_fc1 = self.fc1(x)\n",
    "        x_fc2 = self.fc2(x)\n",
    "        x = self.silu(x_fc1) * x_fc2\n",
    "        return self.fc3(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "qcD8LSHNhBRW",
   "metadata": {
    "id": "qcD8LSHNhBRW"
   },
   "source": [
    "请注意，我们在上述代码中添加了 `dtype=cfg[\"dtype\"]` 设置，使我们能够在 **低精度格式** 下直接加载模型，以减少内存占用（相比于先以 32 位精度实例化后再转换）。\n",
    "\n",
    "此外，我们设置了 `bias=False`，因为 Llama **不使用偏置单元（bias units）**。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6b7bf4f-99d0-42c1-807c-5074d2cc1949",
   "metadata": {
    "id": "f6b7bf4f-99d0-42c1-807c-5074d2cc1949"
   },
   "source": [
    "&nbsp;\n",
    "## 1.4 实现 RoPE（旋转位置编码）"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3487a6f-0373-49d8-b2eb-f8ee05d42884",
   "metadata": {
    "id": "d3487a6f-0373-49d8-b2eb-f8ee05d42884"
   },
   "source": [
    "在 GPT 模型中，位置嵌入（positional embeddings）实现如下：\n",
    "\n",
    "```python\n",
    "self.pos_emb = nn.Embedding(cfg[\"context_length\"], cfg[\"emb_dim\"])\n",
    "```\n",
    "\n",
    "与传统的 **绝对位置编码** 不同，Llama 采用 **旋转位置编码（RoPE）**，能够同时捕捉 **绝对位置** 和 **相对位置** 信息。\n",
    "\n",
    "RoPE 相关论文：[RoFormer: Enhanced Transformer with Rotary Position Embedding (2021)](https://arxiv.org/abs/2104.09864)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a34180fb-448f-44e9-a244-0c736051687b",
   "metadata": {
    "id": "a34180fb-448f-44e9-a244-0c736051687b"
   },
   "outputs": [],
   "source": [
    "def precompute_rope_params(head_dim, theta_base=10_000, context_length=4096):\n",
    "    assert head_dim % 2 == 0, \"Embedding dimension must be even\"\n",
    "\n",
    "    # Compute the inverse frequencies\n",
    "    inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2)[: (head_dim // 2)].float() / head_dim))\n",
    "\n",
    "    # Generate position indices\n",
    "    positions = torch.arange(context_length)\n",
    "\n",
    "    # Compute the angles\n",
    "    angles = positions[:, None] * inv_freq[None, :]  # Shape: (context_length, head_dim // 2)\n",
    "\n",
    "    # Expand angles to match the head_dim\n",
    "    angles = torch.cat([angles, angles], dim=1)  # Shape: (context_length, head_dim)\n",
    "\n",
    "    # Precompute sine and cosine\n",
    "    cos = torch.cos(angles)\n",
    "    sin = torch.sin(angles)\n",
    "\n",
    "    return cos, sin\n",
    "\n",
    "def compute_rope(x, cos, sin):\n",
    "    # x: (batch_size, num_heads, seq_len, head_dim)\n",
    "    batch_size, num_heads, seq_len, head_dim = x.shape\n",
    "    assert head_dim % 2 == 0, \"Head dimension must be even\"\n",
    "\n",
    "    # Split x into first half and second half\n",
    "    x1 = x[..., : head_dim // 2]  # First half\n",
    "    x2 = x[..., head_dim // 2 :]  # Second half\n",
    "\n",
    "    # Adjust sin and cos shapes\n",
    "    cos = cos[:seq_len, :].unsqueeze(0).unsqueeze(0)  # Shape: (1, 1, seq_len, head_dim)\n",
    "    sin = sin[:seq_len, :].unsqueeze(0).unsqueeze(0)\n",
    "\n",
    "    # Apply the rotary transformation\n",
    "    rotated = torch.cat((-x2, x1), dim=-1)\n",
    "    x_rotated = (x * cos) + (rotated * sin)\n",
    "\n",
    "    return x_rotated.to(dtype=x.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e841b8e-75aa-49db-b1a7-d5c2116dc299",
   "metadata": {
    "id": "8e841b8e-75aa-49db-b1a7-d5c2116dc299"
   },
   "source": [
    "- 以下是将 **旋转位置编码（RoPE, Rotary Positional Embedding）** 应用于 `q` 和 `k` 张量的示例：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8c89f022-7167-4001-8c21-8e012878733f",
   "metadata": {
    "id": "8c89f022-7167-4001-8c21-8e012878733f"
   },
   "outputs": [],
   "source": [
    "# Settings\n",
    "batch_size = 2\n",
    "context_len = 5\n",
    "num_heads = 4\n",
    "head_dim = 16\n",
    "\n",
    "# Instantiate RoPE parameters\n",
    "cos, sin = precompute_rope_params(head_dim=head_dim, context_length=context_len)\n",
    "\n",
    "# Dummy query and key tensors\n",
    "torch.manual_seed(123)\n",
    "queries = torch.randn(batch_size, num_heads, context_len, head_dim)\n",
    "keys = torch.randn(batch_size, num_heads, context_len, head_dim)\n",
    "\n",
    "# Apply rotary position embeddings\n",
    "queries_rot = compute_rope(queries, cos, sin)\n",
    "keys_rot = compute_rope(keys, cos, sin)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f78127b0-dda2-4c5a-98dd-bae8f5fe8297",
   "metadata": {
    "id": "f78127b0-dda2-4c5a-98dd-bae8f5fe8297"
   },
   "source": [
    "&nbsp;\n",
    "## 1.5 在 `MultiHeadAttention` 模块中添加 RoPE"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "RnmSHROLhhR3",
   "metadata": {
    "id": "RnmSHROLhhR3"
   },
   "source": [
    "- 需要注意的是，GPT 将 **位置编码（Positional Embedding）** 应用于输入，而 LLaMA 则在 **自注意力机制（Self-Attention）** 中对查询（`q`）和键（`k`）向量执行旋转变换。  \n",
    "- 在此，我们修改 `MultiHeadAttention` 类，以适配 **旋转位置编码（RoPE, Rotary Positional Embedding）**。  \n",
    "- 另外，我们移除 `qkv_bias` 选项，并将 `bias=False` 作为固定设置。  \n",
    "- 此外，我们添加了 `dtype` 设置，以便后续能够使用更低精度的数值格式实例化模型。  \n",
    "- **提示**：由于 `TransformerBlock`（在下一部分）是完全重复的，我们可以优化代码，使每个 `MultiHeadAttention` 模块仅初始化一次缓冲区。然而，我们仍将 **预计算的 RoPE 参数** 添加到 `MultiHeadAttention` 类中，使其可以作为独立模块运行。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d81a441e-0b79-4a8b-8291-ea7f55d58c84",
   "metadata": {
    "id": "d81a441e-0b79-4a8b-8291-ea7f55d58c84"
   },
   "outputs": [],
   "source": [
    "#####################################\n",
    "# Chapter 3\n",
    "#####################################\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, d_in, d_out, context_length, num_heads, dtype=None):  # ,dropout, num_heads, qkv_bias=False):\n",
    "        super().__init__()\n",
    "        assert d_out % num_heads == 0, \"d_out must be divisible by n_heads\"\n",
    "\n",
    "        self.d_out = d_out\n",
    "        self.num_heads = num_heads\n",
    "        self.head_dim = d_out // num_heads  # Reduce the projection dim to match desired output dim\n",
    "\n",
    "        ################################### NEW ###################################\n",
    "        # Set bias=False and dtype=dtype for all linear layers below\n",
    "        ###########################################################################\n",
    "        self.W_query = nn.Linear(d_in, d_out, bias=False, dtype=dtype)\n",
    "        self.W_key = nn.Linear(d_in, d_out, bias=False, dtype=dtype)\n",
    "        self.W_value = nn.Linear(d_in, d_out, bias=False, dtype=dtype)\n",
    "        self.out_proj = nn.Linear(d_out, d_out, bias=False, dtype=dtype)  # Linear layer to combine head outputs\n",
    "        # self.dropout = nn.Dropout(dropout)\n",
    "        self.register_buffer(\"mask\", torch.triu(torch.ones(context_length, context_length), diagonal=1))\n",
    "\n",
    "        ################################### NEW ###################################\n",
    "        cos, sin = precompute_rope_params(head_dim=self.head_dim, context_length=context_length)\n",
    "        self.register_buffer(\"cos\", cos)\n",
    "        self.register_buffer(\"sin\", sin)\n",
    "        ###########################################################################\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        b, num_tokens, d_in = x.shape\n",
    "\n",
    "        keys = self.W_key(x)  # Shape: (b, num_tokens, d_out)\n",
    "        queries = self.W_query(x)\n",
    "        values = self.W_value(x)\n",
    "\n",
    "        # We implicitly split the matrix by adding a `num_heads` dimension\n",
    "        # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)\n",
    "        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "        values = values.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)\n",
    "\n",
    "        # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)\n",
    "        keys = keys.transpose(1, 2)\n",
    "        queries = queries.transpose(1, 2)\n",
    "        values = values.transpose(1, 2)\n",
    "\n",
    "        ################################### NEW ###################################\n",
    "        keys = compute_rope(keys, self.cos, self.sin)\n",
    "        queries = compute_rope(queries, self.cos, self.sin)\n",
    "        ###########################################################################\n",
    "\n",
    "        # Compute scaled dot-product attention (aka self-attention) with a causal mask\n",
    "        attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head\n",
    "\n",
    "        # Original mask truncated to the number of tokens and converted to boolean\n",
    "        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]\n",
    "\n",
    "        # Use the mask to fill attention scores\n",
    "        attn_scores.masked_fill_(mask_bool, -torch.inf)\n",
    "\n",
    "        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
    "        # attn_weights = self.dropout(attn_weights)\n",
    "\n",
    "        # Shape: (b, num_tokens, num_heads, head_dim)\n",
    "        context_vec = (attn_weights @ values).transpose(1, 2)\n",
    "\n",
    "        # Combine heads, where self.d_out = self.num_heads * self.head_dim\n",
    "        context_vec = context_vec.reshape(b, num_tokens, self.d_out)\n",
    "        context_vec = self.out_proj(context_vec)  # optional projection\n",
    "\n",
    "        return context_vec"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "-lt9SfnVioB3",
   "metadata": {
    "id": "-lt9SfnVioB3"
   },
   "source": [
    "- 下面是一个使用 `MultiHeadAttention` 模块处理示例输入的示例："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "03f15755-0083-483f-963b-99b599651638",
   "metadata": {
    "id": "03f15755-0083-483f-963b-99b599651638"
   },
   "outputs": [],
   "source": [
    "# Settings\n",
    "batch_size = 1\n",
    "context_len = 100\n",
    "max_context_len = 4096\n",
    "embed_dim = 128\n",
    "num_heads = 4\n",
    "\n",
    "\n",
    "example_batch = torch.randn((batch_size, context_len, embed_dim))\n",
    "\n",
    "mha = MultiHeadAttention(\n",
    "    d_in=embed_dim,\n",
    "    d_out=embed_dim,\n",
    "    context_length=max_context_len,\n",
    "    num_heads=num_heads\n",
    ")\n",
    "\n",
    "mha(example_batch)\n",
    "\n",
    "del mha  # delete to free up memory"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5a1a272-a038-4b8f-aaaa-f4b241e7f23f",
   "metadata": {
    "id": "e5a1a272-a038-4b8f-aaaa-f4b241e7f23f"
   },
   "source": [
    "&nbsp;\n",
    "## 1.6 更新TransformerBlock 模块"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "255f70ac-9c2e-4328-8af7-1c298b8d4a18",
   "metadata": {
    "id": "255f70ac-9c2e-4328-8af7-1c298b8d4a18"
   },
   "source": [
    "- 在此阶段，大部分的核心工作已经完成；现在我们可以更新 `TransformerBlock` 来使用我们在上面实现的代码。  \n",
    "- 具体而言，我们需要：\n",
    "  - 将 **LayerNorm** 替换为 **RMSNorm**  \n",
    "  - 移除 **dropout**  \n",
    "  - 移除 **`qkv_bias` 设置**  \n",
    "  - 添加 **`dtype` 设置**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2e110721-bf2b-42b3-989a-1635b1658af0",
   "metadata": {
    "id": "2e110721-bf2b-42b3-989a-1635b1658af0"
   },
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.att = MultiHeadAttention(\n",
    "            d_in=cfg[\"emb_dim\"],\n",
    "            d_out=cfg[\"emb_dim\"],\n",
    "            context_length=cfg[\"context_length\"],\n",
    "            num_heads=cfg[\"n_heads\"],\n",
    "            dtype=cfg[\"dtype\"]  # NEW\n",
    "            # dropout=cfg[\"drop_rate\"],\n",
    "            # qkv_bias=cfg[\"qkv_bias\"]\n",
    "        )\n",
    "        self.ff = FeedForward(cfg)\n",
    "\n",
    "        ################################### NEW ###################################\n",
    "        # self.norm1 = LayerNorm(cfg[\"emb_dim\"])\n",
    "        # self.norm2 = LayerNorm(cfg[\"emb_dim\"])\n",
    "        self.norm1 = RMSNorm(cfg[\"emb_dim\"])\n",
    "        self.norm2 = RMSNorm(cfg[\"emb_dim\"])\n",
    "        ###########################################################################\n",
    "\n",
    "        # self.drop_shortcut = nn.Dropout(cfg[\"drop_rate\"])\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Shortcut connection for attention block\n",
    "        shortcut = x\n",
    "        x = self.norm1(x)\n",
    "        x = self.att(x)   # Shape [batch_size, num_tokens, emb_size]\n",
    "        # x = self.drop_shortcut(x)\n",
    "        x = x + shortcut  # Add the original input back\n",
    "\n",
    "        # Shortcut connection for feed-forward block\n",
    "        shortcut = x\n",
    "        x = self.norm2(x)\n",
    "        x = self.ff(x)\n",
    "        # x = self.drop_shortcut(x)\n",
    "        x = x + shortcut  # Add the original input back\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ada953bc-e2c0-4432-a32d-3f7efa3f6e0f",
   "metadata": {
    "id": "ada953bc-e2c0-4432-a32d-3f7efa3f6e0f"
   },
   "source": [
    "&nbsp;\n",
    "## 1.7 更新模型类"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba5d991a-559b-47be-96f4-31b881ab2da8",
   "metadata": {
    "id": "ba5d991a-559b-47be-96f4-31b881ab2da8"
   },
   "source": [
    "- 正如你可能还记得的，在 [第 5 章](../01_main-chapter-code/ch05.ipynb) 中，`TransformerBlock` 是主模型中重复使用的基础模块。  \n",
    "- 我们的 LLaMA 模型已接近完成；现在只需更新 `TransformerBlock` 相关的模型代码。  \n",
    "- 具体而言，我们需要：\n",
    "  - 移除 **绝对位置编码**，因为我们已经使用了 **RoPE 编码**  \n",
    "  - 将 **LayerNorm** 替换为 **RMSNorm**  \n",
    "  - 移除 **dropout**  \n",
    "  - 添加 **`dtype` 设置**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "cf8240fe-5d7f-4e7e-b1ac-e0755aab5e79",
   "metadata": {
    "id": "cf8240fe-5d7f-4e7e-b1ac-e0755aab5e79"
   },
   "outputs": [],
   "source": [
    "# class GPTModel(nn.Module):\n",
    "class Llama2Model(nn.Module):\n",
    "    def __init__(self, cfg):\n",
    "        super().__init__()\n",
    "        self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"], dtype=cfg[\"dtype\"])\n",
    "        # self.pos_emb = nn.Embedding(cfg[\"context_length\"], cfg[\"emb_dim\"])\n",
    "        # self.drop_emb = nn.Dropout(cfg[\"drop_rate\"])\n",
    "\n",
    "        self.trf_blocks = nn.Sequential(\n",
    "            *[TransformerBlock(cfg) for _ in range(cfg[\"n_layers\"])])\n",
    "\n",
    "        ################################### NEW ###################################\n",
    "        # self.final_norm = LayerNorm(cfg[\"emb_dim\"])\n",
    "        self.final_norm = RMSNorm(cfg[\"emb_dim\"])\n",
    "        ###########################################################################\n",
    "        self.out_head = nn.Linear(cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False, dtype=cfg[\"dtype\"])\n",
    "\n",
    "    def forward(self, in_idx):\n",
    "        # batch_size, seq_len = in_idx.shape\n",
    "        tok_embeds = self.tok_emb(in_idx)\n",
    "        # pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))\n",
    "        x = tok_embeds  # + pos_embeds  # Shape [batch_size, num_tokens, emb_size]\n",
    "        # x = self.drop_emb(x)\n",
    "        x = self.trf_blocks(x)\n",
    "        x = self.final_norm(x)\n",
    "        logits = self.out_head(x)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bc94940-aaeb-45b9-9399-3a69b8043e60",
   "metadata": {
    "id": "4bc94940-aaeb-45b9-9399-3a69b8043e60"
   },
   "source": [
    "&nbsp;\n",
    "## 2. 模型初始化"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bG--zY-Ljj1f",
   "metadata": {
    "id": "bG--zY-Ljj1f"
   },
   "source": [
    "- 现在模型代码已经完成，我们可以开始初始化它。  \n",
    "- 在 [第 5 章](../01_main-chapter-code/ch05.ipynb) 中，我们使用以下配置文件来指定 **124M 参数规模的 GPT 模型**："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4b7428df-3d02-4ccd-97b5-a629bdabbe8f",
   "metadata": {
    "id": "4b7428df-3d02-4ccd-97b5-a629bdabbe8f"
   },
   "outputs": [],
   "source": [
    "GPT_CONFIG_124M = {\n",
    "    \"vocab_size\": 50257,     # Vocabulary size\n",
    "    \"context_length\": 1024,  # Context length\n",
    "    \"emb_dim\": 768,          # Embedding dimension\n",
    "    \"n_heads\": 12,           # Number of attention heads\n",
    "    \"n_layers\": 12,          # Number of layers\n",
    "    \"drop_rate\": 0.1,        # Dropout rate\n",
    "    \"qkv_bias\": False        # Query-Key-Value bias\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bVi8uiBjw2T",
   "metadata": {
    "id": "8bVi8uiBjw2T"
   },
   "source": [
    "- 作为参考，下面是 **1.5B 参数规模的 GPT 模型** 配置文件： "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "tAOojV_mkEnd",
   "metadata": {
    "id": "tAOojV_mkEnd"
   },
   "outputs": [],
   "source": [
    "GPT_CONFIG_1558M = {\n",
    "    \"vocab_size\": 50257,     # Vocabulary size\n",
    "    \"context_length\": 1024,  # Context length\n",
    "    \"emb_dim\": 1600,         # Embedding dimension\n",
    "    \"n_heads\": 25,           # Number of attention heads\n",
    "    \"n_layers\": 48,          # Number of layers\n",
    "    \"drop_rate\": 0.1,        # Dropout rate\n",
    "    \"qkv_bias\": False        # Query-Key-Value bias\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "HoGGRAGykQTE",
   "metadata": {
    "id": "HoGGRAGykQTE"
   },
   "source": [
    "- 同样，我们可以定义 **7B 参数规模的 LLaMA 2** 配置文件（为简化起见，这里忽略更大规模的模型）："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e0564727-2d35-4f0c-b0fc-cde1e9134a18",
   "metadata": {
    "id": "e0564727-2d35-4f0c-b0fc-cde1e9134a18"
   },
   "outputs": [],
   "source": [
    "LLAMA2_CONFIG_7B = {\n",
    "    \"vocab_size\": 32000,     # Vocabulary size\n",
    "    \"context_length\": 4096,  # Context length\n",
    "    \"emb_dim\": 4096,         # Embedding dimension\n",
    "    \"n_heads\": 32,           # Number of attention heads\n",
    "    \"n_layers\": 32,          # Number of layers\n",
    "    \"hidden_dim\": 11008,     # NEW: Size of the intermediate dimension in FeedForward\n",
    "    \"dtype\": torch.bfloat16  # NEW: Lower-precision dtype to reduce memory usage\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "FAP7fiBzkaBz",
   "metadata": {
    "id": "FAP7fiBzkaBz"
   },
   "source": [
    "- 使用这些设置，我们现在可以初始化 **LLaMA 2 7B 模型**（请注意，这需要约 **26GB** 的内存）。  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7004d785-ac9a-4df5-8760-6807fc604686",
   "metadata": {
    "id": "7004d785-ac9a-4df5-8760-6807fc604686"
   },
   "outputs": [],
   "source": [
    "model = Llama2Model(LLAMA2_CONFIG_7B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "6079f747-8f20-4c6b-8d38-7156f1101729",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6079f747-8f20-4c6b-8d38-7156f1101729",
    "outputId": "0a0eb34b-1a21-4c11-804f-b40007bda5a3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of parameters: 6,738,415,616\n"
     ]
    }
   ],
   "source": [
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Total number of parameters: {total_params:,}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "Bx14NtzWk2wj",
   "metadata": {
    "id": "Bx14NtzWk2wj"
   },
   "source": [
    "- 如上所示，该模型包含 **67 亿个参数**（通常四舍五入后称为 **7B 模型**）。  \n",
    "- 此外，我们可以使用下面的代码计算该模型的 **内存需求**："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "0df1c79e-27a7-4b0f-ba4e-167fe107125a",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0df1c79e-27a7-4b0f-ba4e-167fe107125a",
    "outputId": "11ced939-556d-4511-d5c0-10a94ed3df32"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "float32 (PyTorch default): 52.33 GB\n",
      "bfloat16: 26.17 GB\n"
     ]
    }
   ],
   "source": [
    "def model_memory_size(model, input_dtype=torch.float32):\n",
    "    total_params = 0\n",
    "    total_grads = 0\n",
    "    for param in model.parameters():\n",
    "        # Calculate total number of elements per parameter\n",
    "        param_size = param.numel()\n",
    "        total_params += param_size\n",
    "        # Check if gradients are stored for this parameter\n",
    "        if param.requires_grad:\n",
    "            total_grads += param_size\n",
    "\n",
    "    # Calculate buffer size (non-parameters that require memory)\n",
    "    total_buffers = sum(buf.numel() for buf in model.buffers())\n",
    "\n",
    "    # Size in bytes = (Number of elements) * (Size of each element in bytes)\n",
    "    # We assume parameters and gradients are stored in the same type as input dtype\n",
    "    element_size = torch.tensor(0, dtype=input_dtype).element_size()\n",
    "    total_memory_bytes = (total_params + total_grads + total_buffers) * element_size\n",
    "\n",
    "    # Convert bytes to gigabytes\n",
    "    total_memory_gb = total_memory_bytes / (1024**3)\n",
    "\n",
    "    return total_memory_gb\n",
    "\n",
    "print(f\"float32 (PyTorch default): {model_memory_size(model, input_dtype=torch.float32):.2f} GB\")\n",
    "print(f\"bfloat16: {model_memory_size(model, input_dtype=torch.bfloat16):.2f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "zudd-5PulKFL",
   "metadata": {
    "id": "zudd-5PulKFL"
   },
   "source": [
    "- 最后，如果适用，我们还可以将模型部署到 **NVIDIA GPU** 或 **Apple Silicon GPU**："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a4c50e19-1402-45b6-8ccd-9077b2ba836d",
   "metadata": {
    "id": "a4c50e19-1402-45b6-8ccd-9077b2ba836d"
   },
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = torch.device(\"mps\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dc64a06-27dc-46ec-9e6d-1700a8227d34",
   "metadata": {
    "id": "5dc64a06-27dc-46ec-9e6d-1700a8227d34"
   },
   "source": [
    "&nbsp;\n",
    "## 3. 加载tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eb30f0c-6144-4bed-87d9-6b2bac377005",
   "metadata": {
    "id": "0eb30f0c-6144-4bed-87d9-6b2bac377005"
   },
   "source": [
    "- 在本节中，我们将加载 **模型的分词器（Tokenizer）**。  \n",
    "- **LLaMA 2** 采用 **Google 的 [SentencePiece](https://github.com/google/sentencepiece) 分词器**，而不是 **OpenAI 的 [Tiktoken](https://github.com/openai/tiktoken)**（但 **LLaMA 3** 采用了 **Tiktoken**）。  \n",
    "- **Meta AI** 在 **Hugging Face Hub** 上分享了 **LLaMA 2 原始模型权重**及 **分词器词汇表**。  \n",
    "- 我们将从 **Hub** 下载 **分词器词汇表**，并将其加载到 **SentencePiece** 中。  \n",
    "- 取消注释并运行以下代码，以安装所需的库："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "768989ea-dc60-4dc8-ae84-cbb3fd224422",
   "metadata": {
    "id": "768989ea-dc60-4dc8-ae84-cbb3fd224422"
   },
   "outputs": [],
   "source": [
    "# !pip install huggingface_hub sentencepiece"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "KbnlzsbYmJU6",
   "metadata": {
    "id": "KbnlzsbYmJU6"
   },
   "source": [
    "- 请注意，**Meta AI** 要求您在下载文件之前 **接受 LLaMA 2 许可协议**。  \n",
    "  - 为此，您需要创建一个 **Hugging Face Hub 账号**，然后访问 [meta-llama/Llama-2-7b](https://huggingface.co/meta-llama/Llama-2-7b) 仓库并接受条款。  \n",
    "- 接下来，您需要创建一个 **访问令牌（Access Token）**。  \n",
    "  - 要生成具有 **READ 权限** 的访问令牌，请点击 **右上角的个人头像**，然后选择 **\"Settings\"（设置）**。  \n",
    "\n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/gpt-to-llama/settings.webp?1\" width=\"300px\">\n",
    "\n",
    "- 然后，创建并复制 **访问令牌**，以便在接下来的代码单元中使用。  \n",
    "\n",
    "<img src=\"https://sebastianraschka.com/images/LLMs-from-scratch-images/bonus/gpt-to-llama/access-token.webp?1\" width=\"600px\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "3357a230-b678-4691-a238-257ee4e80185",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3357a230-b678-4691-a238-257ee4e80185",
    "outputId": "768ed6af-ce14-40bc-ca18-117b4b448269"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The token has not been saved to the git credentials helper. Pass `add_to_git_credential=True` in this function directly or `--add-to-git-credential` if using via `huggingface-cli` if you want to set the git credential as well.\n",
      "Token is valid (permission: read).\n",
      "Your token has been saved to /root/.cache/huggingface/token\n",
      "Login successful\n"
     ]
    }
   ],
   "source": [
    "from huggingface_hub import login\n",
    "import json\n",
    "\n",
    "with open(\"config.json\", \"r\") as config_file:\n",
    "    config = json.load(config_file)\n",
    "    access_token = config[\"HF_ACCESS_TOKEN\"]\n",
    "\n",
    "login(token=access_token)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "IxGh6ZYQo0VN",
   "metadata": {
    "id": "IxGh6ZYQo0VN"
   },
   "source": [
    "- 通过 **访问令牌（Access Token）** 登录后，我们的账户将通过验证，确认已接受 **LLaMA 2 许可协议**，然后即可下载 **分词器词汇表**。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "69714ea8-b9b8-4687-8392-f3abb8f93a32",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 153,
     "referenced_widgets": [
      "e6c75a6aa7b942fe84160e286e3acb3d",
      "08f0bf9459bd425498a5cb236f9d4a72",
      "10251d6f724e43788c41d4b7879cbfd3",
      "53a973c0853b44418698136bd04df039",
      "bdb071e7145a4007ae01599333e72612",
      "6b1821a7f4574e3aba09c1e410cc81e4",
      "8c2873eaec3445888ad3d54ad7387950",
      "0c8f7044966e4207b12352503c67dcbb",
      "8b5951213c9e4798a258146d61d02d11",
      "2c05df3f91e64df7b33905b1065a76f7",
      "742ae5487f2648fcae7ca8e22c7f8db9"
     ]
    },
    "id": "69714ea8-b9b8-4687-8392-f3abb8f93a32",
    "outputId": "c230fec9-5c71-4a41-90ab-8a34d114ea01"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:89: UserWarning: \n",
      "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
      "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
      "You will be able to reuse this secret in all of your notebooks.\n",
      "Please note that authentication is recommended but still optional to access public models or datasets.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e6c75a6aa7b942fe84160e286e3acb3d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.model:   0%|          | 0.00/500k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "tokenizer_file = hf_hub_download(\n",
    "    repo_id=\"meta-llama/Llama-2-7b\",\n",
    "    filename=\"tokenizer.model\",\n",
    "    local_dir=\"Llama-2-7b\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "gp7iQ8cXAJLv",
   "metadata": {
    "id": "gp7iQ8cXAJLv"
   },
   "source": [
    "- 为了提供更直观的分词器接口，我们定义一个 **`LlamaTokenizer` 包装类**："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "Ef4WxhjOBOOc",
   "metadata": {
    "id": "Ef4WxhjOBOOc"
   },
   "outputs": [],
   "source": [
    "import sentencepiece as spm\n",
    "\n",
    "\n",
    "class LlamaTokenizer:\n",
    "    def __init__(self, tokenizer_file):\n",
    "        sp = spm.SentencePieceProcessor()\n",
    "        sp.load(tokenizer_file)\n",
    "        self.tokenizer = sp\n",
    "\n",
    "    def encode(self, text):\n",
    "        return self.tokenizer.encode_as_ids(text)\n",
    "\n",
    "    def decode(self, ids):\n",
    "        return self.tokenizer.decode_pieces(ids)\n",
    "\n",
    "\n",
    "tokenizer = LlamaTokenizer(tokenizer_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "NVhmFeX3pT_M",
   "metadata": {
    "id": "NVhmFeX3pT_M"
   },
   "source": [
    "- 现在，我们可以使用 `generate` 函数，让 **LLaMA 2** 模型生成新的文本。  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e0a2b5cd-6cba-4d72-b8ff-04d8315d483e",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "e0a2b5cd-6cba-4d72-b8ff-04d8315d483e",
    "outputId": "acd5065d-8900-4ba8-ef85-968365f3a0cb"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Every effort movesαllRadius deletingpretcc否']; future eer napulate lackус während inter DES издаSchéon로жа Bass differencespadxsnu ;; ctx始\n"
     ]
    }
   ],
   "source": [
    "from previous_chapters import generate, text_to_token_ids, token_ids_to_text\n",
    "\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "token_ids = generate(\n",
    "    model=model,\n",
    "    idx=text_to_token_ids(\"Every effort moves\", tokenizer).to(device),\n",
    "    max_new_tokens=30,\n",
    "    context_size=LLAMA2_CONFIG_7B[\"context_length\"],\n",
    "    top_k=1,\n",
    "    temperature=0.\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93WTtAA5paYV",
   "metadata": {
    "id": "93WTtAA5paYV"
   },
   "source": [
    "- 当然，如上所示，生成的文本目前是 **无意义的**，因为我们尚未训练 **LLaMA 2** 模型。  \n",
    "- 在下一节，我们不会自行训练模型（这将花费 **数万到数十万美元**），而是直接 **加载 Meta AI 提供的预训练权重**。  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f63cc248-1d27-4eb6-aa50-173b436652f8",
   "metadata": {
    "id": "f63cc248-1d27-4eb6-aa50-173b436652f8"
   },
   "source": [
    "&nbsp;\n",
    "## 4.加载预训练权重"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aKeN7rUfqZMI",
   "metadata": {
    "id": "aKeN7rUfqZMI"
   },
   "source": [
    "- 下面我们加载 **[\"meta-llama/Llama-2-7b\"](https://huggingface.co/meta-llama/Llama-2-7b)** 预训练基础模型，该模型在微调之前仅用于 **文本补全**。  \n",
    "- 或者，您可以通过修改下一代码单元中的字符串，加载 **指令微调并对齐的 [\"meta-llama/Llama-2-7b-chat\"](https://huggingface.co/meta-llama/Llama-2-7b-chat) 模型**。  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "5fa9c06c-7a53-4b4d-9ce4-acc027322ee4",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "66e777955e8748df878f118f07f38dab",
      "da89ae3ea4d2474e98f64ada608f3cea",
      "93e6da39c25f4edfaa72056c89df1f7f",
      "b628603e4cb0405398c916587ee96756",
      "93bedcb9245e44a0a1eb7e4155070f66",
      "0723f467d37b4904819a8bb33ebda10f",
      "e54928776bc649339002adced63738b0",
      "d8e0f42068af4cb094e2f115f76e06e0",
      "0a939565b6e94f08bee0a66e0f9827d4",
      "a5fedbb7ec2e43d99711bb4cd84b9486",
      "0c186f6539714d8eab023969ce47c500"
     ]
    },
    "id": "5fa9c06c-7a53-4b4d-9ce4-acc027322ee4",
    "outputId": "0d8942cc-e5e2-4e77-ec41-1ac7bec7d94f"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "66e777955e8748df878f118f07f38dab",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "consolidated.00.pth:   0%|          | 0.00/13.5G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "weights_file = hf_hub_download(\n",
    "   repo_id=\"meta-llama/Llama-2-7b\",\n",
    "   filename=\"consolidated.00.pth\",\n",
    "   local_dir=\"Llama-2-7b\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "e67cca5c-ba4b-4be5-85c7-fdceae8a5701",
   "metadata": {
    "id": "e67cca5c-ba4b-4be5-85c7-fdceae8a5701"
   },
   "outputs": [],
   "source": [
    "weights = torch.load(weights_file, weights_only=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "-15SJ7btq2zE",
   "metadata": {
    "id": "-15SJ7btq2zE"
   },
   "source": [
    "- `weights` 变量包含以下张量（仅为简洁起见，展示前 **15 个**）：  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ee26bd0b-fea9-4924-97f7-409c14f28e49",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ee26bd0b-fea9-4924-97f7-409c14f28e49",
    "outputId": "fa83d38a-bb41-4cb2-d3c7-c573bfe1f8a4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['tok_embeddings.weight',\n",
       " 'norm.weight',\n",
       " 'output.weight',\n",
       " 'layers.0.attention.wq.weight',\n",
       " 'layers.0.attention.wk.weight',\n",
       " 'layers.0.attention.wv.weight',\n",
       " 'layers.0.attention.wo.weight',\n",
       " 'layers.0.feed_forward.w1.weight',\n",
       " 'layers.0.feed_forward.w2.weight',\n",
       " 'layers.0.feed_forward.w3.weight',\n",
       " 'layers.0.attention_norm.weight',\n",
       " 'layers.0.ffn_norm.weight',\n",
       " 'layers.1.attention.wq.weight',\n",
       " 'layers.1.attention.wk.weight',\n",
       " 'layers.1.attention.wv.weight']"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(weights.keys())[:15]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "UeeSpnunrDFB",
   "metadata": {
    "id": "UeeSpnunrDFB"
   },
   "source": [
    "- 下面的函数基于 **[第 5 章](../01_main-chapter-code/ch05.ipynb) 中的 `load_weights_into_gpt`**，用于将 **预训练权重加载到 LLaMA 2 模型** 中。  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3820e2a7-4f26-41bc-953b-f3879b0aff65",
   "metadata": {
    "id": "3820e2a7-4f26-41bc-953b-f3879b0aff65"
   },
   "outputs": [],
   "source": [
    "def assign(left, right):\n",
    "    if left.shape != right.shape:\n",
    "        raise ValueError(f\"Shape mismatch. Left: {left.shape}, Right: {right.shape}\")\n",
    "\n",
    "    if isinstance(right, torch.Tensor):\n",
    "        return torch.nn.Parameter(right.clone().detach())\n",
    "    else:\n",
    "        return torch.nn.Parameter(torch.tensor(right))\n",
    "\n",
    "\n",
    "def load_weights_into_llama(model, param_config, params):\n",
    "    model.tok_emb.weight = assign(model.tok_emb.weight, params[\"tok_embeddings.weight\"])\n",
    "\n",
    "    for l in range(param_config[\"n_layers\"]):\n",
    "\n",
    "        # Load attention weights\n",
    "        model.trf_blocks[l].att.W_query.weight = assign(\n",
    "            model.trf_blocks[l].att.W_query.weight,\n",
    "            params[f\"layers.{l}.attention.wq.weight\"]\n",
    "        )\n",
    "        model.trf_blocks[l].att.W_key.weight = assign(\n",
    "            model.trf_blocks[l].att.W_key.weight,\n",
    "            params[f\"layers.{l}.attention.wk.weight\"]\n",
    "        )\n",
    "        model.trf_blocks[l].att.W_value.weight = assign(\n",
    "            model.trf_blocks[l].att.W_value.weight,\n",
    "            params[f\"layers.{l}.attention.wv.weight\"]\n",
    "        )\n",
    "        model.trf_blocks[l].att.out_proj.weight = assign(\n",
    "            model.trf_blocks[l].att.out_proj.weight,\n",
    "            params[f\"layers.{l}.attention.wo.weight\"]\n",
    "        )\n",
    "        model.trf_blocks[l].norm1.weight = assign(\n",
    "            model.trf_blocks[l].norm1.weight,\n",
    "            params[f\"layers.{l}.attention_norm.weight\"]\n",
    "        )\n",
    "\n",
    "        # Load FeedForward weights\n",
    "        model.trf_blocks[l].ff.fc1.weight = assign(\n",
    "            model.trf_blocks[l].ff.fc1.weight,\n",
    "            params[f\"layers.{l}.feed_forward.w1.weight\"]\n",
    "        )\n",
    "        # For some reason w2 and w3 are provided in the wrong order in the weights file\n",
    "        model.trf_blocks[l].ff.fc2.weight = assign(\n",
    "            model.trf_blocks[l].ff.fc2.weight,\n",
    "            params[f\"layers.{l}.feed_forward.w3.weight\"]\n",
    "        )\n",
    "        model.trf_blocks[l].ff.fc3.weight = assign(\n",
    "            model.trf_blocks[l].ff.fc3.weight,\n",
    "            params[f\"layers.{l}.feed_forward.w2.weight\"]\n",
    "        )\n",
    "        model.trf_blocks[l].norm2.weight = assign(\n",
    "            model.trf_blocks[l].norm2.weight,\n",
    "            params[f\"layers.{l}.ffn_norm.weight\"]\n",
    "        )\n",
    "\n",
    "    # Load output layer weights\n",
    "    model.final_norm.weight = assign(model.final_norm.weight, params[\"norm.weight\"])\n",
    "    model.out_head.weight = assign(model.out_head.weight, params[\"output.weight\"])\n",
    "\n",
    "\n",
    "load_weights_into_llama(model, LLAMA2_CONFIG_7B, weights)\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "TDuv_Us2rNvk",
   "metadata": {
    "id": "TDuv_Us2rNvk"
   },
   "source": [
    "- 现在，我们已经准备好使用该模型进行 **文本生成**。  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "240987e8-a023-462e-9376-9edfb27559ec",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "240987e8-a023-462e-9376-9edfb27559ec",
    "outputId": "044f24b3-4018-4860-834d-6c2731b9e47c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " Every effort has been made to ensure that the information contained in this website is accurate and up to date and correct at the time of publication\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(123)\n",
    "\n",
    "token_ids = generate(\n",
    "    model=model,\n",
    "    idx=text_to_token_ids(\"Every effort\", tokenizer).to(device),\n",
    "    max_new_tokens=25,\n",
    "    context_size=LLAMA2_CONFIG_7B[\"context_length\"],\n",
    "    top_k=1,\n",
    "    temperature=0.\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d72ed949-b6c0-4966-922f-eb0da732c404",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "## 5. 使用指令微调模型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "akyo7WNyF_YL",
   "metadata": {
    "id": "akyo7WNyF_YL"
   },
   "source": [
    "- 如前所述，我们之前使用的是 **预训练基础模型**。  \n",
    "- 如果您希望使用 **能够遵循指令的模型**，请改用 `\"meta-llama/Llama-2-7b-chat\"`，如下所示："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "nbvAV7vaz6yc",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 101,
     "referenced_widgets": [
      "3b2448a60f5f4ba5b2c686037c8ecd78",
      "60c5932944f24f5fad1d8da89c8e5ae9",
      "aa31aed1b8854a4281fd7e81c60e1205",
      "d4acf06c2414412f8f2fb4f48981c954",
      "693d69251d3d48219c084af17b54b851",
      "ff36d28c55dd4db3a0f76a87640fdfe2",
      "71c49ef820494d5f8908a3daf39f0755",
      "525dc406534f4369b11208816f8fd0d7",
      "865f39213a7341b68f2fe73caaf801b1",
      "eaf4c0231b6d4993b2f8e9e63d8b6921",
      "a11edf3b018e42c88a63a515cf7fe478"
     ]
    },
    "id": "nbvAV7vaz6yc",
    "outputId": "724f5508-d976-4e31-b3d7-95fa65b2c1e8"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3b2448a60f5f4ba5b2c686037c8ecd78",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "consolidated.00.pth:   0%|          | 0.00/13.5G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text:\n",
      " What do llamas eat?\n",
      "Llamas and alpacas are herbivores, which means they eat grasses, leaves, grass\n"
     ]
    }
   ],
   "source": [
    "del model  # to free up memory\n",
    "\n",
    "weights_file = hf_hub_download(\n",
    "   repo_id=\"meta-llama/Llama-2-7b-chat\",\n",
    "   filename=\"consolidated.00.pth\",\n",
    "   local_dir=\"Llama-2-7b-chat\"\n",
    ")\n",
    "\n",
    "model = Llama2Model(LLAMA2_CONFIG_7B)\n",
    "load_weights_into_llama(model, LLAMA2_CONFIG_7B, weights)\n",
    "model.to(device);\n",
    "\n",
    "torch.manual_seed(123)\n",
    "\n",
    "token_ids = generate(\n",
    "    model=model,\n",
    "    idx=text_to_token_ids(\"What do llamas eat?\", tokenizer).to(device),\n",
    "    max_new_tokens=25,\n",
    "    context_size=LLAMA2_CONFIG_7B[\"context_length\"],\n",
    "    top_k=1,\n",
    "    temperature=0.\n",
    ")\n",
    "\n",
    "print(\"Output text:\\n\", token_ids_to_text(token_ids, tokenizer))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f693da1-a07c-4e1d-af5a-c3923525f1e2",
   "metadata": {},
   "source": [
    "&nbsp;\n",
    "# 展望"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fae93739-ca12-46ba-8ca7-7c07c59f669b",
   "metadata": {},
   "source": [
    "- 本笔记本将原始 **GPT-2** 架构转换为 **LLaMA 2** 模型。  \n",
    "- 如果您对如何将 **LLaMA 2** 转换为 **LLaMA 3**、**LLaMA 3.1** 和 **LLaMA 3.2** 感兴趣，请查看 [converting-llama2-to-llama3.ipynb](converting-llama2-to-llama3.ipynb) 笔记本。  "
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "0723f467d37b4904819a8bb33ebda10f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "08f0bf9459bd425498a5cb236f9d4a72": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_6b1821a7f4574e3aba09c1e410cc81e4",
      "placeholder": "​",
      "style": "IPY_MODEL_8c2873eaec3445888ad3d54ad7387950",
      "value": "tokenizer.model: 100%"
     }
    },
    "0a939565b6e94f08bee0a66e0f9827d4": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "0c186f6539714d8eab023969ce47c500": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "0c8f7044966e4207b12352503c67dcbb": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "10251d6f724e43788c41d4b7879cbfd3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_0c8f7044966e4207b12352503c67dcbb",
      "max": 499723,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_8b5951213c9e4798a258146d61d02d11",
      "value": 499723
     }
    },
    "2c05df3f91e64df7b33905b1065a76f7": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "3b2448a60f5f4ba5b2c686037c8ecd78": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_60c5932944f24f5fad1d8da89c8e5ae9",
       "IPY_MODEL_aa31aed1b8854a4281fd7e81c60e1205",
       "IPY_MODEL_d4acf06c2414412f8f2fb4f48981c954"
      ],
      "layout": "IPY_MODEL_693d69251d3d48219c084af17b54b851"
     }
    },
    "525dc406534f4369b11208816f8fd0d7": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "53a973c0853b44418698136bd04df039": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_2c05df3f91e64df7b33905b1065a76f7",
      "placeholder": "​",
      "style": "IPY_MODEL_742ae5487f2648fcae7ca8e22c7f8db9",
      "value": " 500k/500k [00:00&lt;00:00, 3.39MB/s]"
     }
    },
    "60c5932944f24f5fad1d8da89c8e5ae9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_ff36d28c55dd4db3a0f76a87640fdfe2",
      "placeholder": "​",
      "style": "IPY_MODEL_71c49ef820494d5f8908a3daf39f0755",
      "value": "consolidated.00.pth: 100%"
     }
    },
    "66e777955e8748df878f118f07f38dab": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_da89ae3ea4d2474e98f64ada608f3cea",
       "IPY_MODEL_93e6da39c25f4edfaa72056c89df1f7f",
       "IPY_MODEL_b628603e4cb0405398c916587ee96756"
      ],
      "layout": "IPY_MODEL_93bedcb9245e44a0a1eb7e4155070f66"
     }
    },
    "693d69251d3d48219c084af17b54b851": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "6b1821a7f4574e3aba09c1e410cc81e4": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "71c49ef820494d5f8908a3daf39f0755": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "742ae5487f2648fcae7ca8e22c7f8db9": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "865f39213a7341b68f2fe73caaf801b1": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "8b5951213c9e4798a258146d61d02d11": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "8c2873eaec3445888ad3d54ad7387950": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "93bedcb9245e44a0a1eb7e4155070f66": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "93e6da39c25f4edfaa72056c89df1f7f": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d8e0f42068af4cb094e2f115f76e06e0",
      "max": 13476925163,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_0a939565b6e94f08bee0a66e0f9827d4",
      "value": 13476925163
     }
    },
    "a11edf3b018e42c88a63a515cf7fe478": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "a5fedbb7ec2e43d99711bb4cd84b9486": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "aa31aed1b8854a4281fd7e81c60e1205": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_525dc406534f4369b11208816f8fd0d7",
      "max": 13476925163,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_865f39213a7341b68f2fe73caaf801b1",
      "value": 13476925163
     }
    },
    "b628603e4cb0405398c916587ee96756": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a5fedbb7ec2e43d99711bb4cd84b9486",
      "placeholder": "​",
      "style": "IPY_MODEL_0c186f6539714d8eab023969ce47c500",
      "value": " 13.5G/13.5G [01:40&lt;00:00, 111MB/s]"
     }
    },
    "bdb071e7145a4007ae01599333e72612": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "d4acf06c2414412f8f2fb4f48981c954": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_eaf4c0231b6d4993b2f8e9e63d8b6921",
      "placeholder": "​",
      "style": "IPY_MODEL_a11edf3b018e42c88a63a515cf7fe478",
      "value": " 13.5G/13.5G [02:52&lt;00:00, 81.1MB/s]"
     }
    },
    "d8e0f42068af4cb094e2f115f76e06e0": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "da89ae3ea4d2474e98f64ada608f3cea": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_0723f467d37b4904819a8bb33ebda10f",
      "placeholder": "​",
      "style": "IPY_MODEL_e54928776bc649339002adced63738b0",
      "value": "consolidated.00.pth: 100%"
     }
    },
    "e54928776bc649339002adced63738b0": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "e6c75a6aa7b942fe84160e286e3acb3d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_08f0bf9459bd425498a5cb236f9d4a72",
       "IPY_MODEL_10251d6f724e43788c41d4b7879cbfd3",
       "IPY_MODEL_53a973c0853b44418698136bd04df039"
      ],
      "layout": "IPY_MODEL_bdb071e7145a4007ae01599333e72612"
     }
    },
    "eaf4c0231b6d4993b2f8e9e63d8b6921": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "ff36d28c55dd4db3a0f76a87640fdfe2": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
